'''
Yuyang
202403
'''
import dv_processing as dv
import cv2
import numpy as np
#import torch
import os
import glob
import argparse
from tqdm import tqdm

#from davis_utils.buffer import FrameBuffer, EventBuffer
from metavision_core.event_io import EventsIterator

#torch.manual_seed(0)
#torch.cuda.manual_seed(0)
np.random.seed(0)



def parse_argument():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ## dir params
    parser.add_argument('--raw_file_dir', type=str,
                        default='H:\Datasets/val/normal',
                        help='The path of a raw dir')
    parser.add_argument('--save_dir', type=str,
                        default='H:\Datasets/val/frames',
                        help='Output directory for time surface images')
    parser.add_argument('--time_interval_us', type=int, default=50000,
                        help='')

    parser.add_argument('--show', default=True, help='if show detection and tracking process')
    return parser


def get_reader(file_path, time_interval_us):
    assert os.path.exists(file_path), 'The file \'{}\' is not exist'.format(file_path)
    mv_iterator = EventsIterator(input_path=file_path, delta_t=time_interval_us)

    return mv_iterator



if __name__ == '__main__':
    ## Get params
    args, _ = parse_argument().parse_known_args(None)
    print(args)


    #raw_path_list = glob.glob(os.path.join(args.raw_file_dir, '*.raw'))
    raw_path_list = []
    for root, dirs, files in os.walk(args.raw_file_dir):
        for file in files:
            if file.endswith('.raw'):
                raw_path_list.append(os.path.join(root, file))

    
    print(raw_path_list)


    for raw_path in tqdm(raw_path_list):
        assert os.path.exists(raw_path)
        base_name = os.path.basename(raw_path).split('.')[0]
        print(base_name)

        ev_save_path = os.path.join(args.save_dir, base_name)
        ##ev_save_path = os.path.join(args.save_dir, base_name, 'event')
        if not os.path.exists(ev_save_path):
            os.makedirs(ev_save_path)

        reader = get_reader(raw_path, args.time_interval_us)
        height, width = reader.get_size()


        ## 1. save image
        pass


        ## 2. save event
        ev_cnt = 0
        for evs in reader:
            min_t, max_t = np.min(evs['t']), np.max(evs['t'])
            event_image = np.zeros((height, width))
            np.add.at(event_image, (evs['y'], evs['x']), 1)
            cv2.imwrite('{}/{:06d}_{}_{}.png'.format(ev_save_path, ev_cnt, min_t, max_t), (event_image * 60).astype(np.uint8))
            ev_cnt += 1


        # file_start_time, file_end_time = reader.getTimeRange()
        # cur_time = file_start_time
        # ev_cnt = 0
        # while cur_time < file_end_time:
        #     event_store = reader.getEventsTimeRange(cur_time, cur_time+args.time_interval_us)
        #     events, _, _ = EventBuffer.store_to_ndarray(event_store)
        #     cur_time += args.time_interval_us
        #
        #     event_image = np.zeros((height, width))
        #     np.add.at(event_image, (events[:, 1], events[:, 0]), 1)
        #
        #     cv2.imwrite('{}/{:06d}_{}.png'.format(ev_save_path, ev_cnt, cur_time), (event_image*80).astype(np.uint8))
        #     ev_cnt += 1





















